package com.alibaba.cloud.ai.example.graph.basic_chatbot;

import lombok.extern.slf4j.Slf4j;
import org.bsc.langgraph4j.CompileConfig;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.boot.CommandLineRunner;
import org.springframework.stereotype.Controller;
import reactor.core.publisher.Flux;

import java.util.Map;

@Slf4j
@Controller
public class BasicChatConsoleController implements CommandLineRunner {
    private final BasicChatbot basicChatbot;

    public BasicChatConsoleController(BasicChatbot basicChatbot) {
        this.basicChatbot = basicChatbot;
    }

    /**
     * Executes the command-line interface to demonstrate a Spring Boot application.
     * This method logs a welcome message, constructs a graph using an agent executor,
     * compiles it into a workflow, invokes the workflow with a specific input,
     * and then logs the final result.
     *
     * @param args Command line arguments (Unused in this context)
     * @throws Exception If any error occurs during execution
     */
    @Override
    public void run(String... args) throws Exception {

        log.info("Welcome to the Spring Boot CLI application!");
        log.info("==============start==============");
        CompileConfig compileConfig = CompileConfig.builder()
                .checkpointSaver(new BlankSaver())
                .build();

        var graph = basicChatbot.graphBuilder().build();

        var workflow = graph.compile(compileConfig);

        var result = workflow.invoke(Map.of(BasicChatbot.State.INPUT, "what is the capital of China ?"));

        var finish = result.flatMap(BasicChatbot.State::agentOutcome)
                .map(BasicChatbot.Outcome::finish)
                .orElseThrow();

        Map<String, Object> returnValues = finish.returnValues();

        //获取Flux<ChatResponse>类型返回，并进行输出
        Flux<ChatResponse> chatResponseFlux = (Flux<ChatResponse>) returnValues.get("returnValues");

        StringBuilder stringBuilder = new StringBuilder();
        chatResponseFlux.flatMap(chatResponse -> {
            String response = chatResponse.getResult().getOutput().getContent();
            log.info("chatResponse: {}", response);
            stringBuilder.append(response);
            return Flux.empty();
        }).blockLast();
        log.info("==============end==============");
        log.info(stringBuilder.toString());
    }
}